import sys
import pandas as pd
from train_test import (
    mobile_price_run,
    cifar_image_run,
    jaychou_lyrics_run
)
from utils import (
    print_help,
    Logger,
    LOG_FMT
)


def main(cmd: str, log: bool = True) -> None:
    """
    Main entry point for executing different machine learning tasks.

    Args:
        cmd: Command string specifying which task to run.
        log: Boolean flag to enable/disable logging (default: True).
    """
    print_help()  # Display help information for available commands
    # Mobile price prediction task
    if cmd == 'mobile_price':
        # Load mobile price dataset
        mobile_price = pd.read_csv('./data/mobile_price.csv')
        # Initialize logger if enabled
        logger = Logger(
            log_path='./log/',
            log_name=f'mobile_price_{LOG_FMT}.log',
            level='INFO'
        ).get_logger() if log else None
        # Execute mobile price prediction pipeline
        mobile_price_run(mobile_price, logger)

    # CIFAR-10 image classification task    
    elif cmd == 'cifar10_image':
        # Initialize logger if enabled
        logger = Logger(
            log_path='./log/',
            log_name=f'cifar_image_{LOG_FMT}.log',
            level='INFO'
        ).get_logger() if log else None
        # Execute CIFAR-10 image classification pipeline
        cifar_image_run('./data/', logger)

    # Jay Chou lyrics generation task
    elif cmd == 'jaychou_lyrics':
        # Initialize logger if enabled
        logger = Logger(
            log_path='./log/',
            log_name=f'jaychou_lyrics_{LOG_FMT}.log',
            level='INFO'
        ).get_logger() if log else None
        # Execute lyrics generation with specified parameters
        jaychou_lyrics_run(
            file_path='./data/jaychou_lyrics.txt',
            prefix='分手',
            max_len=50,
            logger=logger
        )


if __name__ == '__main__':
    command = sys.argv[1]
    log = True if len(sys.argv) == 3 and sys.argv[2] in ['-l', '--log'] else False
    main(cmd=command, log=log)
